#!/usr/bin/env python

from __future__ import with_statement

import planningproblem
import parserbasic
import codeparser
import feature
import stateparser
import actionparser
import goalparser
import planner
import output
import pyoutput
import re
import sys
import os
import shutil
#argparse needs Python 2.7
#import argparse
import getopt

def parse_problem(fn):
    print "Loading",fn,"..."
    with open(fn,'r') as fin:
        code = fin.readlines()

    linenos,sections = parserbasic.split_sections(code, 
                ['CODE:', 'INITIAL:', 'STATECONSTRAINTS:',
                'FEATURES:', re.compile('ACTION (.+):'), 'GOAL:', 
                'PLANNER:', 'OUTPUT:', 'EXEC:'])

    problem = planningproblem.PlanningProblem()

    problem.bind_globals()

    #parse code setup and set the environment
    codeParser = codeparser.CodeParser(problem.env)
    if 'CODE:' in linenos:
        codeParser.parse(sections['CODE:'],linenos['CODE:'][0])

    #parse initial state and state constraints
    stateParser = stateparser.StateParser(problem.env)
    if 'INITIAL:' in linenos:
        stateParser.parse(sections['INITIAL:'],linenos['INITIAL:'][0])
    if 'STATECONSTRAINTS:' in linenos:
        stateParser.parse_constraints(sections['STATECONSTRAINTS:'],linenos['STATECONSTRAINTS:'][0])
    problem.setInitial(stateParser.space(),stateParser.initial())
    
    featureParser = feature.FeatureParser(problem.env)
    if 'FEATURES:' in linenos:
        featureParser.parse(sections['FEATURES:'],linenos['FEATURES:'][0])
    problem.setStaticFeatures(featureParser.staticFeatures)
    problem.setDynamicFeatures(featureParser.dynamicFeatures)

    #parse goal test
    if 'GOAL:' in linenos:
        goalParser = goalparser.GoalParser(problem.env)
        goal = goalParser.parse(sections['GOAL:'], linenos['GOAL:'][0])
        problem.setGoalTest(goal)

    for (k,v) in linenos.iteritems():
        if k.startswith('ACTION'):
            name = k[7:-1].strip()
            actionParser = actionparser.ActionParser(problem.env,problem.stateSpace.names)
            action = actionParser.parse(name,sections[k],v[0])
            problem.addAction(action)

    #TODO: what if the CODE: section requires the ability to access the
    #initial state or call the actions?  Would this be a cyclic dependency?
    problem.bind_builtins()
    
    print "Done."

    planneropts = {}
    if 'PLANNER:' in linenos:
        for line in sections['PLANNER:']:
            if len(line.strip())==0: continue
            lineParts = [l.strip() for l in line.split('=')]
            planneropts[lineParts[0]] = eval(lineParts[1],problem.env)

    outputParser = output.OutputParser(problem.env)
    outputter = None
    if 'OUTPUT:' in linenos:
        outputter = outputParser.parse(sections['OUTPUT:'],linenos['OUTPUT:'][0])

    if 'EXEC:' in linenos:
        #HACK: override planning code using problem code
        local = {}
        settings = planner.PlanSettings(**planneropts)
        local['__planner__'] = planner.Planner(problem,settings)
        local['__outputter__'] = outputter
        local['__initial__'] = problem.initialState
        code = codeparser.CodeObject(sections['EXEC:'],'EXEC section',linenos['EXEC:'][0])
        code.execute(problem.env,local)
    
    return problem,planneropts,outputter

def usage():
    print("USAGE: "+sys.argv[0]+" [options] file")
    print("    Plans a solution for the specified PyTAMP file")
    print("OPTIONS:")
    print("""    -h, --help: prints this message
    -i n, --iters=n: uses n outer iterations of the planner
    -j, --jacobiantest: tests whether the jacobians are properly defined
    -r, --redirect: redirects output to stdout.txt""")

if __name__ == "__main__":
    #parser = argparse.ArgumentParser(description='PyTAMP Planner.')
    #parser.add_argument('file', type=string,
    #            help='A .strips file defining the problem')
    #parser.add_argument('-i','--iters', type=int, target='iters', default=1000,
    #            help='Number of outer planning iterations')
    #parser.add_argument('-r','--redirect',action='store_true',default=False,
    #            help='Redirect stdout to "dump.txt"')
    #parser.add_argument('-d','--dump',action='store_true',default=False,
    #            help='Dump plannig tree to plan.txt"')
    #args = parser.parse_args()
    #
    #problemfile = args.file
    #iters = args.iters
    #jacobianTest = args.jacobianTest
    #dump = args.dump
    #redirect = args.redirect

    #Ugly getopt version
    try:
        optlist, args = getopt.getopt(sys.argv[1:],'hi:jdr',['help','iters=','dump','jacobiantest','redirect'])
    except getopt.GetoptError, err:
        print str(err) 
        usage()
        sys.exit(2)

    iters = 1000
    jacobianTest = False
    dump = False
    redirect = False
    dooutput = True
    for o, a in optlist:
        if o in ("-i","--iters"):
            iters = int(a)
        elif o in ("-h", "--help"):
            usage()
            sys.exit()
        elif o in ("-d", "--dump"):
            dump = True
        elif o in ("-r", "--redirect"):
            redirect = True
        elif o in ("-j","--jacobiantest"):
            jacobianTest = True
        elif o in ("-o", "--output"):
            dooutput = True
        else:
            assert False, "unhandled option"
    problemfile = args[0]

    #Load the problem
    problem,opts,outputter = parse_problem(problemfile)

    #Create an output path in output/[name]/[index]
    (fnbase,ext)=os.path.splitext(os.path.basename(problemfile))
    path = 'output/'+fnbase
    if not os.path.exists('output'):
        os.mkdir('output')
    if not os.path.exists(path):
        os.mkdir(path)
    path = path+'/'
    pathindex = 1
    for i in xrange(1,1000):
        if not os.path.exists(path+str(i)):
            pathindex = i
            break
    try:
        os.mkdir(path+str(i))
    except:
        pass
    path = path+str(i)+'/'

    #Switch the output directory
    shutil.copy(problemfile,path)
    os.chdir(path)
    print "Problem output will be placed in",path

    if redirect:
        print "Planning started..."
        f = open('stdout.txt','w')
        sys.stdout = f

    print "Problem:"
    problem.print_properties()
    print
    #print "Discretized version of problem:"
    #problem.constructDiscrete().print_properties()
    #print

    if jacobianTest:
        print "Performing jacobian test..."
        for a in problem.actionSamplers:
            a.self_test_jacobian(problem.initialState)

    #Make the planner, and 
    settings = planner.PlanSettings(**opts)
    p = planner.Planner(problem,settings)

    res = p.plan(iters)
    p.print_stats()

    if dump:
        print "Dumping tree to plan.txt..."
        p.tree.dump('plan.txt')

    printResult = ''
    if res != None:
        if res == False:
            print "Planning returned failure",
            printResult = 'Planner returned failure'
        else:
            printResult = 'Solved, path saved to path.tgf'
            print "Outputting solution path to path.tgf"
            output.dump_solution_path(res,'path.tgf')

            #evaluate path cost
            sumcost = sum(e.a.cost(problem.env,e.get_head().state) for e in res)
            print "Solution path has cost",sumcost
    else:
        printResult = 'Out of time, best path saved to path.tgf'
        nodes,edges = p.tree.enumerate()
        mindist = problem.goalTest.distance(problem.initialState,problem.env)
        ebest = None
        for e in edges:
            d = problem.goalTest.distance(e.get_tail().state,problem.env)
            if d <= mindist:
                mindist = d
                ebest = e

        print "Outputting best solution so far to path.tgf",
        output.dump_solution_path(p.path_to(ebest),'path.tgf')
        print "Best node goal distance:",mindist
        print "Best node path cost:",p.cost_to(ebest)

    if dooutput:
        #os.system('../../../pyoutput.py '+os.path.basename(problemfile)+' path.tgf')
        plan = pyoutput.LoadedPlan(problem,'path.tgf')
        plan.output(outputter,True)

    if redirect:
        sys.stdout.close()
        sys.stdout = sys.__stdout__
        print printResult
